package com.olive.config;

import org.springframework.security.access.AccessDeniedException;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.context.ReactiveSecurityContextHolder;
import org.springframework.security.core.context.SecurityContextImpl;
import org.springframework.security.web.server.WebFilterExchange;
import org.springframework.security.web.server.authorization.ServerAccessDeniedHandler;
import org.springframework.security.web.server.context.ServerSecurityContextRepository;
import org.springframework.security.web.server.util.matcher.ServerWebExchangeMatcher;
import org.springframework.security.web.server.util.matcher.ServerWebExchangeMatchers;
import org.springframework.util.Assert;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;
import reactor.core.publisher.Mono;

/**
 * 自定义http鉴权filter
 */
public class HttpAuthenticationFilter implements WebFilter {

	private CustomAuthenticationManager customAuthenticationManager;

	private ServerAccessDeniedHandler serverAccessDeniedHandler;

	private ServerWebExchangeMatcher requiresAuthenticationMatcher = ServerWebExchangeMatchers.anyExchange();

	private final ServerSecurityContextRepository securityContextRepository;

	public HttpAuthenticationFilter(CustomAuthenticationManager customAuthenticationManager,
			ServerSecurityContextRepository securityContextRepository) {
		Assert.notNull(customAuthenticationManager, "authenticationManager cannot be null");
		this.customAuthenticationManager = customAuthenticationManager;
		this.securityContextRepository = securityContextRepository;
	}

	@Override
	public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
		return this.requiresAuthenticationMatcher.matches(exchange)
				.filter(ServerWebExchangeMatcher.MatchResult::isMatch)
				.switchIfEmpty(chain.filter(exchange).then(Mono.empty()))
				.flatMap(matchResult -> ReactiveSecurityContextHolder.getContext())
				.flatMap(token -> customAuthenticationManager.authenticate(token.getAuthentication()))
				.flatMap(authentication -> onAuthenticationSuccess(authentication,
						new WebFilterExchange(exchange, chain)))
				.onErrorResume(AccessDeniedException.class, e -> serverAccessDeniedHandler.handle(exchange, e));

	}

	protected Mono<Void> onAuthenticationSuccess(Authentication authentication, WebFilterExchange webFilterExchange) {
		ServerWebExchange exchange = customAuthenticationManager.onSuccess(authentication,
				webFilterExchange.getExchange());
		SecurityContextImpl securityContext = new SecurityContextImpl();
		securityContext.setAuthentication(authentication);
		return this.securityContextRepository.save(exchange, securityContext)
				.then(webFilterExchange.getChain().filter(exchange))
				.contextWrite(ReactiveSecurityContextHolder.withSecurityContext(Mono.just(securityContext)));
	}

	public HttpAuthenticationFilter setServerAccessDeniedHandler(ServerAccessDeniedHandler serverAccessDeniedHandler) {
		this.serverAccessDeniedHandler = serverAccessDeniedHandler;
		return this;
	}

	public HttpAuthenticationFilter setAuthenticationManager(CustomAuthenticationManager customAuthenticationManager) {
		this.customAuthenticationManager = customAuthenticationManager;
		return this;
	}

	public HttpAuthenticationFilter setRequiresAuthenticationMatcher(
			ServerWebExchangeMatcher requiresAuthenticationMatcher) {
		this.requiresAuthenticationMatcher = requiresAuthenticationMatcher;
		return this;
	}
}